package tjs.study.notes.flink.kafka;

import com.alibaba.fastjson.JSONObject;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.AssignerWithPeriodicWatermarks;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.assigners.ProcessingTimeSessionWindows;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.streaming.connectors.kafka.KafkaDeserializationSchema;
import org.apache.flink.util.Collector;
import org.apache.kafka.clients.consumer.ConsumerRecord;

import javax.annotation.Nullable;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Desc: Flink Window & Watermark
 * Created by zhisheng on 2019-05-14
 * blog：http://www.54tianzhisheng.cn/
 * 微信公众号：zhisheng
 */
public class WatermarksTest {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //如果不指定时间的话，默认是 ProcessingTime，但是如果指定为事件事件的话，需要事件中带有时间或者添加时间水印
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

        env.setParallelism(16);// 设置并行数，可以再keyBy后看到效果

        Properties props = new Properties();
        props.put("bootstrap.servers", "localhost:9092");
        props.put("zookeeper.connect", "localhost:2181");
        props.put("group.id", "metric-group");
        props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");  //key 反序列化
        props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.put("auto.offset.reset", "earliest"); // latest earliest
        FlinkKafkaConsumer consumer = new FlinkKafkaConsumer<Student>(
                "student_water1",  //kafka topic
                new StudentKafkaDeserialization(),  // String 序列化
                props);

        consumer.setStartFromEarliest();// 测试代码，每次从头开始消费

        //每x毫秒生成一次水位线（默认200）
//        env.getConfig().setAutoWatermarkInterval(300);

        DataStream<Student> data = env.addSource(consumer);
        // 自定义水位线
        // data.assignTimestampsAndWatermarks(new MyWatermarks(Time.seconds(1)));// fixme 此种写法将无效
        data = data.rebalance()// rebalance后每一个solt都将有数据，默认情况下一个分区一个solt
                .assignTimestampsAndWatermarks(new MyWatermarks(Time.seconds(0)));
        // 使用默认不延期水位线
//        data = data.assignTimestampsAndWatermarks(new AscendingTimestampExtractor<Student>() {
//            @Override
//            public long extractAscendingTimestamp(Student element) {
//                return element.timestamp;
//            }
//        });
        // 使用默认延期水位线
//        data = data.assignTimestampsAndWatermarks(new BoundedOutOfOrdernessTimestampExtractor<Student>(Time.seconds(10)) {
//            @Override
//            public long extractTimestamp(Student element) {
//                return element.timestamp;
//            }
//        });
        data.print();
        data.keyBy(Student::getTimestamp)//keyBy以及之后的计算都采用并行执行
                .window(TumblingEventTimeWindows.of(Time.seconds(3)))
                .apply(
                        //WindowFunction<IN, OUT, KEY, W extends Window>
                        new WindowFunction<Student, String, Long, TimeWindow>() {
                            /**
                             * 处理窗口数据集合
                             * @param s         从keyBy里返回的key值
                             * @param window    窗口类型
                             * @param input     从窗口获取的所有分区数据流
                             * @param out       输出数据流对象
                             * @throws Exception
                             */
                            @Override
                            public void apply(Long age, TimeWindow window, Iterable<Student> input, Collector<String> out) throws Exception {
                                String res = "threadId: " + Thread.currentThread().getId() + " " + getDateStr(window.getStart()) + " -> " + getDateStr(window.getEnd());
                                res = res + ", age: " + age + " , ";
                                AtomicInteger sumAccount = new AtomicInteger();
                                List<String> sumIds = new ArrayList<>();
                                input.forEach((item) -> {
                                    sumIds.add(String.valueOf(item.getTimestamp()));
                                    sumAccount.addAndGet(item.getAcount());
                                });
                                res = res + " sumAccount: " + sumAccount.get() + " sumIds: " + JSONObject.toJSONString(sumIds);
                                out.collect(res);
                            }
                        }).setParallelism(2)
//                .sum("acount")
                .print("acount");

        env.execute("zhisheng —— flink window example");
    }

    public static class MyWatermarks implements AssignerWithPeriodicWatermarks<Student> {
        // 定义x毫秒的容忍间隔时间，即允许数据的最大乱序时间
        private long maxOutofOrderness = 1000;
        // 观察到的最大时间戳
        private long currentMaxTs = Long.MIN_VALUE;

        public MyWatermarks(Time maxOutofOrdernessTime) {
            this.maxOutofOrderness = maxOutofOrdernessTime.toMilliseconds();
        }

        @Nullable
        @Override
        public Watermark getCurrentWatermark() {
            // 生成具有x毫秒容忍度的水位线
            return new Watermark(currentMaxTs == Long.MIN_VALUE ? Long.MIN_VALUE : currentMaxTs - maxOutofOrderness);
        }

        @Override
        public long extractTimestamp(Student element, long previousElementTimestamp) {
            //获取当前记录的时间戳
            long currentTs = element.timestamp;
            // 更新最大的时间戳
            currentMaxTs = Math.max(currentMaxTs, currentTs);
            // 返回记录的时间戳
            return currentTs;
        }
    }

    public static class StudentKafkaDeserialization implements KafkaDeserializationSchema<Student> {
        @Override
        public TypeInformation<Student> getProducedType() {
            return TypeInformation.of(new TypeHint<Student>() {
            });

        }

        @Override
        public boolean isEndOfStream(Student nextElement) {
            return false;
        }

        @Override
        public Student deserialize(ConsumerRecord<byte[], byte[]> record) {
            try {
                Student student = JSONObject.parseObject(record.value(), Student.class);
                if (student.getTimestamp() == null || student.getTimestamp().equals(0L)) {// 优先取自定义的
                    student.setTimestamp(record.timestamp());
                }
                return student;
            } catch (Exception ex) {
                ex.printStackTrace();
            }
            return null;
        }
    }

    public static String getDateStr(long timestamps) {
        return new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date(timestamps));
    }

    public static long getTimestamps(String dateStr) throws ParseException {
        return new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse(dateStr).getTime();
    }
}
